package com.linked.gateway.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.linked.gateway.config.GateWaySecuritySetting;
import com.linked.gateway.config.GatewaySwitchConfig;
import com.linked.gateway.util.RSAUtils;
import com.linked.universal.linkedutil.HashUtil;
import com.linked.universal.linkedutil.RSAUtil;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;


/**
 * @author :dbq
 * @date : 2022/10/27 13:27
 */
@Component
public class ParamSecurityFilter implements GlobalFilter, Ordered {

    private final static Logger logger = LoggerFactory.getLogger(ParamSecurityFilter.class);

    private final ObjectMapper mapper;

    private final GateWaySecuritySetting securitySetting;

    private final GatewaySwitchConfig switchConfig;

    @Autowired
    public ParamSecurityFilter(ObjectMapper mapper, GateWaySecuritySetting securitySetting, GatewaySwitchConfig switchConfig) {
        this.mapper = mapper;
        this.securitySetting = securitySetting;
        this.switchConfig = switchConfig;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        if (!switchConfig.isParamProtectSwitch()) {
            return chain.filter(exchange);
        }


        ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
        Mono<String> modifiedBody = serverRequest.bodyToMono(String.class).flatMap(body -> {
            //可能有些参数中含有转义符
            try {
                byte[] bytes_data = HashUtil.decryptBASE64(body);
                byte[] decrypt_data = RSAUtil.decryptByPrivateKey(bytes_data, securitySetting.getPrivateKey());
                String param_str = new String(decrypt_data);
                if (logger.isInfoEnabled()) {
                    logger.info("解密后的参数为：" + param_str);
                }
                return Mono.just(param_str);
            } catch (Exception e) {
                logger.error("参数解密异常！");
                return Mono.empty();
            }

        });

        //创建BodyInserter修改请求体
        BodyInserter<Mono<String>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove(HttpHeaders.CONTENT_LENGTH);
        //创建CachedBodyOutputMessage并且把请求param加入
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
            /**
             * 一、请求的重写
             * */
            ServerHttpRequestDecorator requestDecorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                @Override
                public Flux<DataBuffer> getBody() {
                    return outputMessage.getBody();
                }
            };
            /**
             * 二、响应的重写
             * */
            ServerHttpResponse originalResponse = exchange.getResponse();
            DataBufferFactory bufferFactory = originalResponse.bufferFactory();
            ServerHttpResponseDecorator responseDecorator = new ServerHttpResponseDecorator(exchange.getResponse()) {

                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

                    //判断是否是有效的返回
                    if (HttpStatus.OK.equals(getStatusCode()) && body instanceof Flux) {
                        String originalResponseContentType = exchange.getAttribute(ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);
                        if (!StringUtils.isEmpty(originalResponseContentType) && originalResponseContentType.contains("application/json")) {
                            Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                            return super.writeWith(fluxBody.buffer().map(dataBuffers -> {//解决返回体分段传输
                                DefaultDataBufferFactory defaultDataBufferFactory = new DefaultDataBufferFactory();
                                DataBuffer join = defaultDataBufferFactory.join(dataBuffers);
                                try {
                                    String JsonStr = join.toString(StandardCharsets.UTF_8);
                                    if (logger.isInfoEnabled()) {
                                        logger.info("返回值加密处理，加密前：" + JsonStr);
                                    }
                                    byte[] json_bytes = RSAUtil.encryptByPublicKey(JsonStr.getBytes(StandardCharsets.UTF_8), securitySetting.getPubicKey());
                                    String json_encrypt = HashUtil.encryptBASE64(json_bytes);
                                    if (logger.isInfoEnabled()) {
                                        logger.info("返回值加密处理，加密后：" + json_encrypt);
                                    }
                                    return bufferFactory.wrap(mapper.writeValueAsBytes(json_encrypt));
                                } catch (Exception e) {
                                    logger.error("返回值加密异常！", e);
                                    return bufferFactory.wrap(new byte[0]);
                                }

                            }));
                        }

                    }
                    return super.writeWith(body);
                }
            };
            /**
             * 三、结果封存完成后返回
             * */
            return chain.filter(exchange.mutate().request(requestDecorator).response(responseDecorator).build());
        }));

    }


    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE + 3;
    }
}
